import cv2
from mmocr.utils.ocr import MMOCR

class SP_Rec:
    def __init__(self, rec_config: object, rec_checkpoint: object) -> object:
        self.recog_model = MMOCR(det=None,
                        recog='SATRN',
                        recog_config=rec_config,
                        recog_ckpt=rec_checkpoint)
        

    def __call__(self, im):
        '''

        :param im: numpy
        :return: string
        '''
        recog_text = self.recog_model.readtext(im)[0]['text']
        return recog_text

    def batch_predict(self,ims):
        if len(ims)==0:
            return None
        results=self.recog_model.readtext(ims,recog_batch_size=len(ims))
        recog_texts=[]
        for res in results:
            recog_texts.append(res['text'])
        return recog_texts


if __name__=='__main__':
    rec_config = '/home/wsl/OCR/mmocr_mlt/configs/textrecog/mlt/satrn_hindi.py'
    rec_checkpoint = '/home/wsl/下载/XianDaoHindiTest/ETRN/ETRN_model.pth'
    model=SP_Rec(rec_config, rec_checkpoint)
    im=cv2.imread('/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/OCR/语种判别MLT19/test/Hindi/5.png')
    res=model(im)
    print(res)